# Reference code: https://github.com/li-plus/DSNet/blob/1804176e2e8b57846beb063667448982273fca89/src/make_dataset.py#L4
# Reference code: https://github.com/e-apostolidis/PGL-SUM/blob/81d0d6d0ee0470775ad759087deebbce1ceffec3/model/configs.py#L10
import cv2
import torch
import os

from pathlib import Path
from tqdm import tqdm

from config import get_config
from generate_summary import generate_summary
from model import set_model
from video_helper import VideoPreprocessor
import io
from utils import stderr_redirector

def get_video_duration(video_path):
    cap = cv2.VideoCapture(video_path)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_rate = int(cap.get(cv2.CAP_PROP_FPS))
    duration = frame_count / frame_rate
    return duration

def pick_frames(video_path, selections):
    cap = cv2.VideoCapture(str(video_path))
    frames = []
    n_frames = 0

    f = io.StringIO()
    with stderr_redirector(f):

        while True:
            if n_frames >= 1800:
                break
            ret, frame = cap.read()

            if not ret:
                break
            
            if selections[n_frames]:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frames.append(frame)
            n_frames += 1
    f.close()
    cap.release()

    return frames

def produce_video(save_path, frames, fps, frame_size):
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(save_path, fourcc, fps, frame_size)
    for frame in frames:
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        out.write(frame)
    out.release()

def main_process(save_path, dir_path, weight_path):
    # Load config
    # config = get_config()

    # create output directory
    out_dir = Path(save_path)
    out_dir.mkdir(parents=True, exist_ok=True)

    device='cuda:0'

    # feature extractor
    video_proc = VideoPreprocessor(
        sample_rate=15,
        device=device
    )

    # search all videos with .mp4 suffix
    video_paths = sorted(Path(dir_path).glob(f'*.mp4'))

    # Load CSTA weights
    model = set_model(
        model_name='GoogleNet_Attention',
        Scale=None,
        Softmax_axis='TD',
        Balance=None,
        Positional_encoding="FPE",
        Positional_encoding_shape='TD',
        Positional_encoding_way='PGL_SUM',
        Dropout_on=True,
        Dropout_ratio=0.6,
        Classifier_on=True,
        CLS_on=True,
        CLS_mix='Final',
        key_value_emb='kv',
        Skip_connection='KC',
        Layernorm=True
    )
    model.load_state_dict(torch.load(weight_path, map_location='cpu'))
    model.to(device)
    model.eval()

    problem_list = []

    # Generate summarized videos
    with torch.no_grad():
        for video_path in tqdm(video_paths,total=len(video_paths),ncols=80,leave=False,desc="Making videos..."):
            if os.path.exists(f'{save_path}/{video_path.stem}.mp4'):
                continue
            try:
            # duration = get_video_duration(video_path)
            # if duration >= 3 * 60:
            #     continue
                print(video_path)
                video_name = video_path.stem
                n_frames, features, cps, pick = video_proc.run(video_path)

                inputs = features.to(device)
                inputs = inputs.unsqueeze(0).expand(3,-1,-1).unsqueeze(0)


                outputs = model(inputs)
                predictions = outputs.squeeze().clone().detach().cpu().numpy().tolist()
                # print(cps.shape, len(predictions), n_frames, pick.shape)
                selections = generate_summary([cps], [predictions], [n_frames], [pick])[0]

                frames = pick_frames(video_path=video_path, selections=selections)
                produce_video(
                    save_path=f'{save_path}/{video_name}.mp4',
                    frames=frames,
                    fps=video_proc.fps,
                    frame_size=(video_proc.frame_width,video_proc.frame_height)
                )
            except Exception as e:
                print(e)
                problem_list.append(video_path)
                continue
    
    with open('./problem_list.txt', 'a') as fin:
        fin.write(str(problem_list))

if __name__=='__main__':
    main_process('./summary_videos/HW006-TOR-Profile005/2024-04-13', '../video-summary/HW006-TOR-Profile005/2024-04-13', './weights/SumMe/split4.pt')
